#define _USE_MATH_DEFINES
#include <cmath>
#include "AlgoRRTStar.h"
#include "AlgoRRTStarPrivate.h"
#include "DebugConfig.h"
#include <algorithm>
#include <iostream>
#include <sstream>
#include <limits>

namespace algo
{

// --- AlgoRRTStar public methods ---
AlgoRRTStar::AlgoRRTStar() : d(new AlgoRRTStarPrivate()) {}
AlgoRRTStar::~AlgoRRTStar() { delete d; }
int AlgoRRTStar::exec(const algobase::TaskConfig& taskConfig, AlgoEvaluatorDB* db, algobase::Paths& paths)
{
    return d->exec(taskConfig, db, paths);
}

void AlgoRRTStar::setDepthRange(double min_depth, double max_depth)
{
    d->setDepthRange(min_depth, max_depth);
}

// --- AlgoRRTStarPrivate implementation ---
AlgoRRTStarPrivate::AlgoRRTStarPrivate() {}
AlgoRRTStarPrivate::~AlgoRRTStarPrivate() {}

void AlgoRRTStarPrivate::setDepthRange(double min_depth, double max_depth)
{
    min_depth_ = min_depth;
    max_depth_ = max_depth;
}

int AlgoRRTStarPrivate::exec(const algobase::TaskConfig& taskConfig, AlgoEvaluatorDB* db, algobase::Paths& paths)
{
    paths.clear();
    setCurrentTaskConfig(&taskConfig);
    db_ = db;

    if (!db_) {
        std::cerr << "DB data is not set for RRT*." << std::endl;
        return -1;
    }

    std::vector<algobase::Polygon> mandatory_zones;
    const auto mandatory_polygons_it = taskConfig.mTypePolygons.find(algobase::AT_Mandatory);
    if (mandatory_polygons_it != taskConfig.mTypePolygons.end()) {
        mandatory_zones = mandatory_polygons_it->second;
    }

    int num_paths_to_generate = 150;
    for (int i = 0; i < num_paths_to_generate; ++i) {
        algobase::Path raw_path = runMultiTargetRRTStar(i, mandatory_zones);
        if (!raw_path.empty()) {
            paths.push_back(raw_path);
        }
    }

    if (paths.empty()) {
        DEBUG_LOG("AlgoRRTStar.log_target", "RRT* could not find any path after " << num_paths_to_generate << " attempts.");
        return -2;
    }

    return 0;
}

algobase::Path AlgoRRTStarPrivate::runMultiTargetRRTStar(int random_seed, const std::vector<algobase::Polygon>& mandatory_zones)
{
    node_storage_.clear();
    node_list_.clear();
    random_generator_ = std::make_unique<std::mt19937>(random_seed);

    algobase::CoordZ start_pos_3d = { current_task_config_->mStartPos[0], current_task_config_->mStartPos[1], getDepth(current_task_config_->mStartPos) };
    auto start_node = std::make_unique<RRTNode>(start_pos_3d);
    node_list_.push_back(start_node.get());
    node_storage_.push_back(std::move(start_node));

    RRTNode* last_connection_node = node_list_.front();

    for (const auto& zone : mandatory_zones) {
        if (!growTreeTowardsTarget({}, &zone, &last_connection_node)) {
            DEBUG_LOG("AlgoRRTStar.log_target", "Failed to connect to a mandatory zone.");
            return {};
        }
    }

    algobase::CoordZ end_pos = { current_task_config_->mEndPos[0], current_task_config_->mEndPos[1], getDepth(current_task_config_->mEndPos) };
    if (!growTreeTowardsTarget(end_pos, nullptr, &last_connection_node)) {
        DEBUG_LOG("AlgoRRTStar.log_target", "Failed to connect to the final goal.");
        return {};
    }

    return reconstructPath(last_connection_node);
}

bool AlgoRRTStarPrivate::growTreeTowardsTarget(const algobase::CoordZ& target_pos, const algobase::Polygon* target_zone, RRTNode** connection_node)
{
    algobase::CoordZ current_target = target_zone ? getPolygonCenter(*target_zone) : target_pos;

    for (int i = 0; i < max_iterations_; ++i) {
        algobase::CoordZ q_rand = sampleRandomPoint(current_target);
        RRTNode* q_near = findNearestNode(q_rand);
        algobase::CoordZ q_new_pos = steer(q_near->position, q_rand);

        if (!isCollisionFree(q_near->position, q_new_pos)) {
            DEBUG_LOG("AlgoRRTStar.log_expansion", "Collision detected for new node. Discarding.");
            continue;
        }

        RRTNode* q_parent = chooseParent(q_near, q_new_pos);
        
        auto q_new_node = std::make_unique<RRTNode>(q_new_pos);
        q_new_node->parent = q_parent;
        q_new_node->cost = q_parent->cost + distance(q_parent->position, q_new_pos);
        RRTNode* q_new_node_ptr = q_new_node.get();
        node_list_.push_back(q_new_node_ptr);
        node_storage_.push_back(std::move(q_new_node));

        rewire(q_new_node_ptr);

        if (target_zone) {
            algobase::Coord q_new_2d = {q_new_pos[0], q_new_pos[1]};
            if (isInside(q_new_2d, *target_zone)) {
                *connection_node = q_new_node_ptr;
                DEBUG_LOG("AlgoRRTStar.log_target", "Successfully connected to a mandatory zone.");
                return true;
            }
        } else {
            if (distance(q_new_pos, target_pos) < goal_threshold_) {
                *connection_node = q_new_node_ptr;
                DEBUG_LOG("AlgoRRTStar.log_target", "Successfully connected to the final goal.");
                return true;
            }
        }
    }
    return false;
}

algobase::CoordZ AlgoRRTStarPrivate::sampleRandomPoint(const algobase::CoordZ& target_pos)
{
    std::uniform_real_distribution<> dis(0.0, 1.0);
    if (dis(*random_generator_) < goal_bias_) {
        return target_pos;
    }

    std::uniform_real_distribution<> dis_lon(current_task_config_->mMinLongitude, current_task_config_->mMaxLongitude);
    std::uniform_real_distribution<> dis_lat(current_task_config_->mMinLatitude, current_task_config_->mMaxLatitude);
    
    algobase::Coord random_coord = {dis_lon(*random_generator_), dis_lat(*random_generator_)};
    
    std::uniform_real_distribution<> dis_depth(min_depth_, max_depth_);
    double random_depth = dis_depth(*random_generator_);
    
    double terrain_depth = getDepth(random_coord);
    if (random_depth < terrain_depth) {
        random_depth = terrain_depth;
    }

    return {random_coord[0], random_coord[1], random_depth};
}

RRTNode* AlgoRRTStarPrivate::findNearestNode(const algobase::CoordZ& point)
{
    RRTNode* nearest = nullptr;
    double min_dist = std::numeric_limits<double>::max();
    for (RRTNode* node : node_list_) {
        double d = distance(node->position, point);
        if (d < min_dist) {
            min_dist = d;
            nearest = node;
        }
    }
    return nearest;
}

algobase::CoordZ AlgoRRTStarPrivate::steer(const algobase::CoordZ& from, const algobase::CoordZ& to)
{
    double d = distance(from, to);
    if (d < step_size_) return to;
    double ratio = step_size_ / d;
    return { from[0] + (to[0] - from[0]) * ratio, from[1] + (to[1] - from[1]) * ratio, from[2] + (to[2] - from[2]) * ratio };
}

std::vector<RRTNode*> AlgoRRTStarPrivate::findNeighbors(const algobase::CoordZ& point)
{
    std::vector<RRTNode*> neighbors;
    for (RRTNode* node : node_list_) {
        if (distance(node->position, point) < search_radius_) {
            neighbors.push_back(node);
        }
    }
    return neighbors;
}

RRTNode* AlgoRRTStarPrivate::chooseParent(RRTNode* q_near, const algobase::CoordZ& q_new_pos)
{
    RRTNode* min_cost_parent = q_near;
    double min_cost = q_near->cost + distance(q_near->position, q_new_pos);
    for (RRTNode* neighbor : findNeighbors(q_new_pos)) {
        double current_cost = neighbor->cost + distance(neighbor->position, q_new_pos);
        if (current_cost < min_cost && isCollisionFree(neighbor->position, q_new_pos)) {
            min_cost = current_cost;
            min_cost_parent = neighbor;
        }
    }
    return min_cost_parent;
}

void AlgoRRTStarPrivate::rewire(RRTNode* q_new)
{
    for (RRTNode* neighbor : findNeighbors(q_new->position)) {
        if (neighbor == q_new->parent) continue;
        double new_cost = q_new->cost + distance(q_new->position, neighbor->position);
        if (new_cost < neighbor->cost && isCollisionFree(q_new->position, neighbor->position)) {
            neighbor->parent = q_new;
            neighbor->cost = new_cost;
        }
    }
}

bool AlgoRRTStarPrivate::isCollisionFree(const algobase::CoordZ& p1, const algobase::CoordZ& p2) const
{
    int steps = 10;
    for (int i = 0; i <= steps; ++i) {
        double t = static_cast<double>(i) / steps;
        algobase::CoordZ p = { p1[0] + t * (p2[0] - p1[0]), p1[1] + t * (p2[1] - p1[1]), p1[2] + t * (p2[2] - p1[2]) };
        algobase::Coord p_2d = {p[0], p[1]};
        if (p[0] < current_task_config_->mMinLongitude || p[0] > current_task_config_->mMaxLongitude ||
            p[1] < current_task_config_->mMinLatitude  || p[1] > current_task_config_->mMaxLatitude) return false;
        if (getDepth(p_2d) > 0 && p[2] < getDepth(p_2d)) return false;
        for (const auto& pair : current_task_config_->mTypePolygons) {
            if (pair.first == algobase::AT_Obstacle || pair.first == algobase::AT_Restricted) {
                for (const auto& polygon : pair.second) if (isInside(p_2d, polygon)) return false;
            }
        }
    }
    return true;
}

algobase::Path AlgoRRTStarPrivate::reconstructPath(RRTNode* end_node)
{
    algobase::Path path;
    RRTNode* current = end_node;
    while (current != nullptr) {
        algobase::PathCoord pc;
        pc.mCoordZ = current->position;
        path.push_back(pc);
        current = current->parent;
    }
    std::reverse(path.begin(), path.end());

    if (path.size() > 1) {
        for (size_t i = 0; i < path.size() - 1; ++i) {
            const auto& p1 = path[i].mCoordZ;
            const auto& p2 = path[i+1].mCoordZ;
            double dx = p2[0] - p1[0];
            double dy = p2[1] - p1[1];
            path[i].mHeading = std::atan2(dy, dx) * 180.0 / M_PI;
            path[i].mSpeed = 5.0;
        }
        path.back().mHeading = path[path.size() - 2].mHeading;
        path.back().mSpeed = path[path.size() - 2].mSpeed;
    } else if (!path.empty()) {
        path.front().mSpeed = 0.0;
    }

    return path;
}

double AlgoRRTStarPrivate::getDepth(const algobase::Coord& coord) const
{
    if (!db_) return 0.0;
    return db_->getData(algobase::EEDT_Depth, current_task_config_->mStartTime, {coord[0], coord[1], 0.0}, 0.0f);
}

bool AlgoRRTStarPrivate::isInside(const algobase::Coord& point, const algobase::Polygon& polygon) const
{
    if (polygon.size() < 3) return false;
    int crossings = 0;
    for (size_t i = 0; i < polygon.size(); ++i) {
        const auto& p1 = polygon[i];
        const auto& p2 = polygon[(i + 1) % polygon.size()];
        if (((p1[1] <= point[1] && point[1] < p2[1]) || (p2[1] <= point[1] && point[1] < p1[1])) &&
            (point[0] < (p2[0] - p1[0]) * (point[1] - p1[1]) / (p2[1] - p1[1]) + p1[0])) {
            crossings++;
        }
    }
    return (crossings % 2) != 0;
}

double AlgoRRTStarPrivate::distance(const algobase::CoordZ& p1, const algobase::CoordZ& p2) const
{
    double lon1 = p1[0], lat1 = p1[1], depth1 = p1[2];
    double lon2 = p2[0], lat2 = p2[1], depth2 = p2[2];
    double R = 6371e3;
    double x = (lon2 - lon1) * M_PI/180.0 * cos((lat1 + lat2)/2.0 * M_PI/180.0);
    double y = (lat2 - lat1) * M_PI/180.0;
    double horiz_dist = sqrt(x*x + y*y) * R;
    double vert_dist = depth2 - depth1;
    return sqrt(horiz_dist*horiz_dist + vert_dist*vert_dist);
}

algobase::CoordZ AlgoRRTStarPrivate::getPolygonCenter(const algobase::Polygon& polygon) const
{
    if (polygon.empty()) return {};
    double avg_lon = 0, avg_lat = 0;
    for(const auto& p : polygon) { avg_lon += p[0]; avg_lat += p[1]; }
    algobase::Coord center_2d = {avg_lon / polygon.size(), avg_lat / polygon.size()};
    return {center_2d[0], center_2d[1], getDepth(center_2d)};
}

} // namespace algo
